# 不用这个文件了，直接用clip处理就行
# -*- coding: utf-8 -*
import argparse
import torch
import torch.nn as nn
import torchvision
import torchvision.models as models
from torch.nn.parallel import DataParallel
from torchvision.datasets import ImageFolder
from torch.autograd import Variable
from torch.optim import Adam
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
from sklearn.decomposition import PCA
import os
import sys

def writetxt(filepath, list, feature_dim):
	"""
	将特征向量存储进txt文件中
	:param filepath: 写入文件地址
	:param list: 提取的图片特征向量
	:param feature_dim: 图片特征向量维度
	:return: 无
	"""
	img_embed_file = open(filepath, 'w')
	for count, feature in enumerate(list):
		for num in feature:
			img_embed_file.write(' ' + str(num))
		img_embed_file.write('\n')
	img_embed_file.close()

def get_mean_std(datapath):
	"""
	计算数据集的均值和方差
	:param datapath: 使用的数据集路径 训练、测试、验证集可以分开处理
	:return: 返回均值mean和方差std
	"""
	means = [0, 0, 0]
	stdevs = [0, 0, 0]
	transform = transforms.Compose([transforms.Resize((224, 224)),
									transforms.CenterCrop((224, 224)),
									transforms.ToTensor(),
									# transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))
									])
	dataset = ImageFolder(datapath, transform)
	num_imgs = len(dataset)
	for j, data in enumerate(dataset):
		print(j)
		img = data[0]
		for i in range(3):
			# 一个通道的均值和方差
			means[i] += img[i, :, :].mean()
			stdevs[i] += img[i, :, :].std()

	means = np.asarray(means) / num_imgs
	stdevs = np.asarray(stdevs) / num_imgs

	print("normMean = {}".format(means))
	print("normstdevs = {}".format(stdevs))

	return means, stdevs

model_name = 'alexnet'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# IMAGE_PATH = os.path.join(sys.path[0], 'data', 'train')
IMAGE_PATH = '/home/hlf/code/code_LSTM+CNN/data_set/hours_2021.6.8_2022.5.31'
img_embed_path = '/home/hlf/code/code_LSTM+CNN/data_set/img_embed.txt'
BATCH = 1
EPOCHS = 1
cuda_avail = torch.cuda.is_available()
# 数据预处理
means, std = get_mean_std(IMAGE_PATH)
transform = transforms.Compose([
	transforms.Resize((224, 224)),
	transforms.ToTensor(),
	transforms.Normalize(means, std)
	# transforms.Normalize([0.4405344, 0.41118264, 0.37459105], [0.23881102, 0.22837733, 0.22760215]) # CELF数据集的均值方差
])

imageSet = ImageFolder(IMAGE_PATH, transform=transform)
data = DataLoader(imageSet, batch_size=BATCH, shuffle=True)

if model_name == 'vgg':
	model = models.vgg19_bn(pretrained=True)
if model_name == 'alexnet':
	model = models.alexnet(pretrained=True)

save_feature = []
for i in range(EPOCHS):
	model.fc = torch.nn.LeakyReLU(0.1)
	model.eval()
	for j, trainData in enumerate(data):
		x_train = trainData[0]
		x_train = x_train.float()

		if cuda_avail:
			x_train = Variable(x_train)

		feature = model(x_train)
		save_feature.append(feature.detach().cpu().numpy().ravel())
		if (j + 1) % 100 == 0:
			print(j, len(data))
save_feature = torch.Tensor(save_feature)
save_feature = save_feature.numpy().tolist()

writetxt(img_embed_path, save_feature, len(save_feature[0]))
